import math
import os, sys
import traceback

#"pip install PyOpenGL"
from OpenGL.GL  import *
from OpenGL.GLU import *

#"pip install pygame-ce"
with open( os.devnull, "w" ) as f:
	oldstdout=sys.stdout; sys.stdout=f; import pygame; sys.stdout=oldstdout

import gl_shader

from config import *



if sys.platform in ("win32","win64"): os.environ["SDL_VIDEO_CENTERED"]="1"

pygame.display.init()
pygame.font.init()



bounds_x_original = list( bounds_x )

##Change around as you like to prezoom on a point (below is this release's screenshot)
#bounds_x = [ 0.4136968978, 0.4139727058 ]
#bounds_y = [ 0.3037255751, 0.3039094471 ]

font = pygame.font.SysFont( None, 12 )

icon=pygame.Surface((32,32)); icon.set_colorkey((0,0,0)); pygame.display.set_icon(icon)
pygame.display.set_caption( "Mandebrot Set Viewer - Agatha Mallett - v.5.03 - 2024" )
pygame.display.set_mode( window_res, pygame.OPENGL|pygame.DOUBLEBUF|pygame.RESIZABLE )



glEnable( GL_BLEND )
glBlendFunc( GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA )

shader_frag = """
#version 120 //OpenGL 2.1, minimum

//#define BLACK_AND_WHITE
#ifndef BLACK_AND_WHITE
	#define HAS_MOD
	#ifdef HAS_MOD
		#extension GL_EXT_gpu_shader4 : enable /*for "%"*/
	#endif
#endif

#if """+str(int(use_double_precision))+"""
	#extension GL_ARB_gpu_shader_fp64 : enable
	#define FLOAT double
	#define FVEC2 dvec2
	#define LOG(X) FLOAT(log(float(X)))
#else
	#define FLOAT float
	#define FVEC2 vec2
	#define LOG(X) log(X)
#endif

#define NUM_COLORS 16
#define MAX_ITERS """+str(MAX_ITERS)+"""
#define AA_SQRT_RES """+str(AA_SQRT_RES)+"""

uniform FVEC2 bounds_x;
uniform FVEC2 bounds_y;
uniform FVEC2 window_res;

vec3 colors[16] = vec3[16](
	vec3(241,233,191) / 255.0,
	vec3(248,201, 95) / 255.0,
	vec3(255,170,  0) / 255.0,
	vec3(204,108,  0) / 255.0,
	vec3(153, 87,  0) / 255.0,
	vec3(106, 52,  3) / 255.0,
	vec3( 66, 30, 15) / 255.0,
	vec3( 25,  7, 26) / 255.0,
	vec3(  9,  1, 47) / 255.0,
	vec3(  4,  4, 73) / 255.0,
	vec3(  0,  7,100) / 255.0,
	vec3( 12, 44,138) / 255.0,
	vec3( 24, 82,177) / 255.0,
	vec3( 57,125,209) / 255.0,
	vec3(134,181,229) / 255.0,
	vec3(211,236,248) / 255.0
);
vec3 get_color( int iter )
{
	#ifdef BLACK_AND_WHITE
		return vec3(255.0);
	#else
		#ifdef HAS_MOD
			iter %= 16;
		#else
			iter -= (iter/16)*16;
		#endif
		return colors[ iter ];
	#endif
}

int   shade_index      ( int iter                 ) { return iter; }
FLOAT shadesmooth_index( int iter, FVEC2 position )
{
	FLOAT normal_iter = FLOAT(iter) - LOG(LOG(length(position))) / LOG(2.0);
	return (normal_iter+0.5) * """+str(1.0/smooth_shading_scale)+""";
}

FVEC2 complex_sq(FVEC2 z)
{
	FVEC2 temp1 = z * z;
	FLOAT temp2 = z.x * z.y;
	return FVEC2( temp1.x-temp1.y, temp2+temp2 );
}
/*FVEC2 complex_mul(FVEC2 a, FVEC2 b)
{
	return FVEC2( a.x*b.x-a.y*b.y, a.y*b.x+a.x*b.y );
}*/

vec3 sample( FVEC2 c )
{
	FVEC2 z = c;
	for ( int i=0; i<MAX_ITERS; ++i )
	{
		//z = complex_add( complex_mul(z,z), c );
		z = complex_sq(z) + c;

		// Want to calculate if length(z)>some radius.  If it is, then we stop.
		// The minimum radius is 2.0, but smooth shading requires more to converge.
		FLOAT z_dot_z = dot( z, z ); // z.x*z.x and z.y*z.y are also cool

		#if """+str(int(not use_smooth_shading))+""" // simple shading
			if ( z_dot_z > 2.0*2.0 ) return get_color( shade_index(i) );
		#else //smooth shading
			if ( z_dot_z > 8.0*8.0 )
			{
				FLOAT index_f = shadesmooth_index( i, z );

				int index_i = int(index_f)
				index_f -= FLOAT(index_i);

				return mix( get_color(index_i),get_color(index_i+1), index_f )
			}
		#endif
	}
	return vec3( 0.0, 0.0, 0.0 );
}
void main()
{
	vec4 color = vec4( 0.0,0.0,0.0, 1.0 );

	FVEC2 win_sc   = FVEC2(bounds_x[1]-bounds_x[0],bounds_y[1]-bounds_y[0]) / window_res;
	FVEC2 win_bias = FVEC2(bounds_x[0],bounds_y[0]);

	FVEC2 coord_pixel00 = FVEC2(gl_FragCoord.xy) - FVEC2(0.5);
	for ( int y=0; y<AA_SQRT_RES; ++y )
	for ( int x=0; x<AA_SQRT_RES; ++x )
	{
		FVEC2 coord_subpixel = ( FVEC2(FLOAT(x),FLOAT(y)) + FVEC2(0.5) ) / FLOAT(AA_SQRT_RES);
		FVEC2 coord_pixel = coord_pixel00 + coord_subpixel;

		FVEC2 c = coord_pixel*win_sc + win_bias;
		color.rgb += sample( c );
	}
	color.rgb /= float( AA_SQRT_RES * AA_SQRT_RES );

	gl_FragData[0] = color;
}
"""
program = gl_shader.Program([ gl_shader.ShaderFrag(shader_frag) ])
#program.save_sources( "program_sources.txt" )



def get_input():
	global bounds_x, bounds_y

	keys_pressed = pygame.key.get_pressed()
	mouse_buttons  = pygame.mouse.get_pressed()
	mouse_position = pygame.mouse.get_pos    ()

	for event in pygame.event.get():
		if   event.type == pygame.QUIT: return False
		elif event.type == pygame.KEYDOWN:
			if   event.key == pygame.K_ESCAPE: return False
			elif event.key == pygame.K_s and (
				keys_pressed[pygame.K_LCTRL] or keys_pressed[pygame.K_RCTRL]
			):
				data = glReadPixels( 0,0,window_res[0],window_res[1], GL_RGB, GL_UNSIGNED_BYTE )
				assert type(data) == bytes
				img = pygame.image.frombytes( data )
				pygame.image.save( img, "screenshot.png" )
		elif event.type == pygame.VIDEORESIZE:
			window_res[0] = event.w
			window_res[1] = event.h
			
	zoom_conv = 0.0
	if mouse_buttons[0]:
		zoom_conv =       zoom_scalar
	if mouse_buttons[2]:
		zoom_conv = 1.0 / zoom_scalar
	if zoom_conv != 0.0:
		fx = float(              mouse_position[0]) / float(window_res[0])
		fy = float(window_res[1]-mouse_position[1]) / float(window_res[1])
		center = [
			fx * (bounds_x[1]-bounds_x[0]) + bounds_x[0],
			fy * (bounds_y[1]-bounds_y[0]) + bounds_y[0]
		]
		
		bounds_x[0] = bounds_x[0]*zoom_conv + center[0]*(1.0-zoom_conv)
		bounds_x[1] = bounds_x[1]*zoom_conv + center[0]*(1.0-zoom_conv)

		bounds_y[0] = bounds_y[0]*zoom_conv + center[1]*(1.0-zoom_conv)
		bounds_y[1] = bounds_y[1]*zoom_conv + center[1]*(1.0-zoom_conv)
		
	return True

def _draw_fractal():
	gl_shader.Program.use( program )
	program.pass_2f( "bounds_x", bounds_x )
	program.pass_2f( "bounds_y", bounds_y )
	program.pass_2f( "window_res", window_res )

	glBegin(GL_QUADS)
	glVertex2f( 0.0          , 0.0           )
	glVertex2f( window_res[0], 0.0           )
	glVertex2f( window_res[0], window_res[1] )
	glVertex2f( 0.0          , window_res[1] )
	glEnd()

	gl_shader.Program.use( None )

def _draw_labels():
	def render_label( text:str ):
		surf = font.render( text, True, (255,255,255,255) )
		surf2 = pygame.Surface(( surf.get_width()+10, 13 )).convert_alpha()
		surf2.fill(( 50,50,50, 192 ))
		surf2.blit( surf, (5,2) )
		return surf2

	def draw_surf( surf, pos ):
		glPushMatrix()
		glRasterPos2i( pos[0], pos[1] )
		glDrawPixels(
			surf.get_width(), surf.get_height(),
			GL_RGBA, GL_UNSIGNED_BYTE, pygame.image.tobytes(surf,"RGBA",True)
		)
		glPopMatrix()

	surf = render_label( f"( {bounds_x[0]}, {bounds_y[1]} i )" )
	draw_surf( surf, (              2,                 window_res[1]-(13+2)) )

	surf = render_label (f"( {bounds_x[0]}, {bounds_y[0]} i )" )
	draw_surf( surf, (              2,                                   2 ) )

	surf = render_label( f"( {bounds_x[1]}, {bounds_y[1]} i )" )
	draw_surf( surf, (window_res[0]-2-surf.get_width(),window_res[1]-(13+2)) )

	surf = render_label( f"( {bounds_x[1]}, {bounds_y[0]} i )" )
	draw_surf( surf, (window_res[0]-2-surf.get_width(),                  2 ) )

	def add_commas( intstr:str ):
		ret = []
		for k,ch in enumerate(reversed(intstr)):
			ret.append(ch)
			if k%3 == 2: ret.append(",")
		if ret[-1] == ",":
			ret.pop()
		return "".join(reversed(ret))

	zoom = int(round(
		float(bounds_x_original[1]-bounds_x_original[0]) / float(bounds_x[1]-bounds_x[0])
	))
	surf = render_label( f"Zoom Factor: {add_commas(str(zoom))}x" )
	draw_surf( surf, (window_res[0]-2-surf.get_width(),17) )

def draw():
	glViewport( 0,0, window_res[0],window_res[1] )

	glMatrixMode(GL_PROJECTION)
	glLoadIdentity()
	glOrtho( 0.0,window_res[0], 0.0,window_res[1], -1.0,1.0 )

	glMatrixMode(GL_MODELVIEW)
	glLoadIdentity()

	glClear( GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT )

	_draw_fractal()

	_draw_labels()

	pygame.display.flip()

def main():
	clock = pygame.time.Clock()
	while True:
		if not get_input(): break
		draw()
		clock.tick(60)

	global program
	del program # Delete before OpenGL context is deleted

	pygame.quit()

if __name__ == "__main__":
	try:
		main()
	except:
		traceback.print_exc()
		pygame.quit()
		input()
